import os
import jieba
from sklearn import metrics
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB


def loadfile(filepath, lables):
    listfile = os.listdir(filepath)
    word_list = []
    lables_list = []
    for file in listfile:
        f = filepath + '/' + file
        word_list.append(cut_word(f))
        lables_list.append(lables)
    return word_list, lables_list


def cut_word(filepath):
    with open(filepath, 'r',encoding='gb18030') as file:
        data = file.read()
        word_txt = ''
        for word in jieba.cut(data):
            word_txt += word+' '
        return word_txt


# 加载训练数据
train_girls_data, train_girls_lables = loadfile('./data/train/女性', '女性')
train_tiyu_data, train_tiyu_lables = loadfile('./data/train/体育', '体育')
train_wenxue_data, train_wenxue_lables = loadfile('./data/train/文学', '文学')
train_xiaoyuan_data, train_xiaoyuan_lables = loadfile('./data/train/校园', '校园')
train_data = train_girls_data + train_tiyu_data + train_wenxue_data + train_xiaoyuan_data
train_lables = train_girls_lables + train_tiyu_lables + train_wenxue_lables + train_xiaoyuan_lables

# 加载测试数据
test_girls_data, test_girls_lables = loadfile('./data/test/女性', '女性')
test_tiyu_data, test_tiyu_lables = loadfile('./data/test/体育', '体育')
test_wenxue_data, test_wenxue_lables = loadfile('./data/test/文学', '文学')
test_xiaoyuan_data, test_xiaoyuan_lables = loadfile('./data/test/校园', '校园')
test_data = test_girls_data + test_tiyu_data + test_wenxue_data + test_xiaoyuan_data
test_lables = test_girls_lables + test_tiyu_lables + test_wenxue_lables + test_xiaoyuan_lables

# 加载停用词
# stop_list = []
# with open('./data/stop/stopword.txt',encoding='utf-8',errors='ignore') as file:
#     for line in file:
#         stop_list.append(line)
stop_words = open('./data/stop/stopword.txt', 'r', encoding='utf-8').read()
stop_words = stop_words.encode('utf-8').decode('utf-8-sig') # 列表头部\ufeff处理
stop_words = stop_words.split('\n') # 根据分隔符分隔

# 计算特征权重
tf = TfidfVectorizer(stop_words=stop_words, max_df=0.5)
train_feature = tf.fit_transform(train_data)
test_feature = tf.transform(test_data)

#训练模型
model = MultinomialNB(alpha = 0.001).fit(train_feature,train_lables)

#模型预测
pre_lables = model.predict(test_feature)

#模型评估
acc = metrics.accuracy_score(test_lables,pre_lables)
print(f'模型准确率是：{acc}')